import torch
import torch.nn as nn
import torch.optim as optim
import random
import numpy as np
from collections import deque

# 定义自动驾驶的神经网络模型
class DriverNetwork(nn.Module):
    def __init__(self, input_size, output_size):
        super(DriverNetwork, self).__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.fc_layers = nn.Sequential(
            nn.Linear(self.input_size + 2, 1024),
            nn.Sigmoid(),
            nn.Linear(1024, 256),
            nn.Sigmoid(),
            nn.Linear(256, self.output_size)
        )

    def forward(self, x):
        x = self.fc_layers(x)
        return x

    def _calculate_hidden_size(self, input_size):
        return ((input_size - 3) // 2 - 4) // 2

# 定义经验回放缓冲区
class ReplayBuffer:
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)

    def push(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):
        return random.sample(self.buffer, batch_size)

    def filt_out(self):
        # 剔除reward小于均值的记录
        rewards = [item[2] for item in self.buffer]
        mean_reward = np.mean(rewards)
        self.buffer = [item for item in self.buffer if item[2] >= mean_reward]

    def __len__(self):
        return len(self.buffer)

# 定义DQN代理
class DqnAgent:
    def __init__(self, input_size, output_size, memory_capacity=100000, learning_rate=1e-3, gamma=0.99, batch_size=64):
        self.dqn = DriverNetwork(input_size, output_size)
        self.target_dqn = DriverNetwork(input_size, output_size)
        self.target_dqn.load_state_dict(self.dqn.state_dict())
        self.target_dqn.eval()
        self.learning_rate = learning_rate
        self.gamma = gamma
        self.batch_size = batch_size
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.dqn.to(self.device)
        self.target_dqn.to(self.device)
        self.optimizer = optim.Adam(self.dqn.parameters(), lr=self.learning_rate)
        self.memory = ReplayBuffer(memory_capacity)
        self.counter=0

    def choose_action(self, state):
        state = torch.tensor(np.array([state]), dtype=torch.float32, device=self.device)
        with torch.no_grad():
            return self.dqn(state).max(1)[1].item()

    def remember(self, state, action, reward, next_state, done):
        self.memory.push(state, action, reward, next_state, done)

    def learn(self):
        if len(self.memory) < self.batch_size:
            return
        transitions = self.memory.sample(self.batch_size)
        states, actions, rewards, next_states, dones = zip(*transitions)
        states = torch.tensor(np.array(states), dtype=torch.float, device=self.device)
        actions = torch.tensor(actions, dtype=torch.int64, device=self.device)
        rewards = torch.tensor(rewards, dtype=torch.float, device=self.device)
        next_states = torch.tensor(np.array(next_states), dtype=torch.float, device=self.device)
        dones = torch.tensor(dones, dtype=torch.uint8, device=self.device)

        q_values = self.dqn(states).gather(1, actions.unsqueeze(1)).squeeze(1)
        with torch.no_grad():
            next_q_values = self.target_dqn(next_states).max(1)[0]
        expected_q_values = rewards + (self.gamma * next_q_values * (1 - dones))

        loss = nn.MSELoss()(q_values, expected_q_values)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        self.counter+=1
        if self.counter%self.batch_size==1:
            self.memory.filt_out()
        if self.counter%300==1:
            self.target_dqn.load_state_dict(self.dqn.state_dict())

    def save(self, filename):
        torch.save(self.target_dqn.state_dict(), filename)

    def load(self, filename):
        self.dqn.load_state_dict(torch.load(filename,weights_only=True))  


if __name__ == "__main__":
    path = "models/DQN_model.pt"
    agent = DqnAgent(
        input_size=41,
        output_size=25,
        memory_capacity=100000,
        batch_size=2048,
        gamma=0.9,
        learning_rate=3e-3,
    )
    agent.load(path)
    print(agent.device)